"""
来自论文[1] Park Woosang, Jung Youeyun, Bang Hyochoong, et al. Robust crater triangle matching algorithm for planetary landing navigation[J]. Journal of Guidance, Control, and Dynamics, 2019, 42(2): 402-410
和论文[2] Meer P., Ramakrishna S., Lenz R. Correspondence of coplanar features through p2-invariant representations[A]. Proceedings of 12th International Conference on Pattern Recognition[C]. Jerusalem, Israel: IEEE Comput. Soc. Press, 1994: 196-200

仅使用论文提供的不变量作为描述子，由于描述子没有顺序，因此只能使用投票法。
"""

import numpy as np
import torch
from utils.ellipse import ellipse_args, radius_ellipse
from .base import TriadPyramidVote
from itertools import combinations
from scipy.optimize import linear_sum_assignment


class ParkPyramidVote(TriadPyramidVote):

    def triad_descriptor(
        self, i, j, k, Q1, Q2, Q3, *args, factor=0.3, **kwargs
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """从陨石坑的参数中计算描述子，C1和C2可以是含有不确定度的椭圆参数"""
        N = Q3.shape[0]
        if len(Q1.shape) == 2:
            Q1 = Q1[None].repeat(N, dim=0)
        if len(Q2.shape) == 2:
            Q2 = Q2[None].repeat(N, dim=0)
        assert Q1.shape[0] == Q2.shape[0] == Q3.shape[0]
        d1 = np.mean(radius_ellipse(Q1), axis=0)
        d2 = np.mean(radius_ellipse(Q2), axis=0)
        d3 = np.mean(radius_ellipse(Q3), axis=0)
        index = np.argsort((d1, d2, d3), axis=0)
        Q = np.array((Q1, Q2, Q3))[index, np.arange(N)]
        x1, y1, a1, b1, _ = ellipse_args(Q[0])
        x2, y2, a2, b2, _ = ellipse_args(Q[1])
        x3, y3, a3, b3, _ = ellipse_args(Q[2])
        d1 = torch.tensor(a1 + b1).to(self.device, dtype=torch.float32)
        d2 = torch.tensor(a2 + b2).to(self.device, dtype=torch.float32)
        d3 = torch.tensor(a3 + b3).to(self.device, dtype=torch.float32)
        Q = torch.tensor(Q).to(self.device, dtype=torch.float32)
        # 三角形的三个顶点坐标
        A = torch.stack(
            (torch.tensor(x1), torch.tensor(y1), torch.ones(x1.shape))
        ).T.to(self.device, dtype=torch.float32)
        B = torch.stack(
            (torch.tensor(x2), torch.tensor(y2), torch.ones(x1.shape))
        ).T.to(self.device, dtype=torch.float32)
        Cs = torch.stack(
            (torch.tensor(x3), torch.tensor(y3), torch.ones(x1.shape))
        ).T.to(self.device, dtype=torch.float32)
        AC = torch.cross(A, Cs, dim=1)
        BC = torch.cross(B, Cs, dim=1)
        AB = torch.cross(A, B, dim=1)
        # 求解交点
        Q = torch.vstack((Q[0], Q[1], Q[2], Q[1], Q[2], Q[0]))
        M = torch.vstack((AB, BC, AC, AB, BC, AC))
        cross = self.solve_cross_points(M, Q)
        # 由于M、Q由六个点组成，因此当一个交点为nan时应当全部舍去
        ## view(-1, 6, 4)分别代表，6个交点，num_tri代表输入的被拉平的三角形个数，但是为了保证
        ## 每个三角形都不存在Nan，因此将拉平结果还原为原始维度，然后判断每个三角形内是否有Nan
        N = cross.shape[0] // 2
        cross = self.judge_inside(
            torch.stack((A, B, Cs), dim=0),
            torch.stack((B, Cs, A), dim=0),
            cross[:N].view(3, -1, 2, 3),
            cross[N:].view(3, -1, 2, 3),
        )
        valid_ind = ~cross.isnan().any(dim=(0, 2, 3))
        # 每行是一个3x6的点列，每三个元素代表一个点
        # 计算不变量
        I = self.g_func(cross.permute(1, 0, 2, 3).flatten(1, 2)).T
        # 计算阈值
        th = 0.1 * I.abs()
        ijk = np.vstack((i, j, k))
        # 按直径从小到大排序
        # ijk = np.take_along_axis(ijk, index, axis=0)
        valid_ind = valid_ind.cpu().numpy()
        return ijk, I, th, np.argsort(index, axis=0), valid_ind

    def identify(self, params, uncertainty, factor, *args, max_try=5, **kwargs):
        """
        输入应该越接近于真实的检测输出越好。因此输入应当是一堆椭圆的拟合参数，输出应当是匹配确定的结果，即椭圆ID的列表。
        Arguments:
            params (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y, 1)
        Returns:
            list : matched ellipse ID
        """
        # 将陨石坑按直径排序
        params = np.array(params)
        elli = ellipse_args(params)
        diameters = elli[2] + elli[3]
        idx = np.argsort(-diameters)
        # 从最大的直径开始，选取陨石坑构建金字塔
        if len(idx) < 4:
            return None
        fail_num = 0
        for i, j in combinations(range(idx.shape[0]), 2):
            k, l = np.array(list(combinations(np.delete(idx, [i, j]), 2))).T
            i = idx[i].repeat(k.shape[0])
            j = idx[j].repeat(k.shape[0])
            # for k, l in combinations(not_used[2:], 2):
            # 批量算出金字塔的三元组
            ijk, I123, th, index, valid = self.triad_descriptor(
                np.concatenate((i, i, j, i)),
                np.concatenate((j, j, k, k)),
                np.concatenate((k, l, l, l)),
                params[np.concatenate((i, i, j, i))],
                params[np.concatenate((j, j, k, k))],
                params[np.concatenate((k, l, l, l))],
                factor,
            )
            # 变成四个三元组，即金字塔的结构
            ijk = ijk.reshape(3, 4, -1)
            I123 = I123.view(I123.shape[0], 4, -1)
            th = th.view(th.shape[0], 4, -1)
            index = index.reshape(3, 4, -1)
            valid = valid.reshape(4, -1).all(axis=0)
            if not valid.any():
                fail_num += 1
                continue
            # 由于上一步是批量计算的，这一步再批量显存要爆了，所以就一行一行的来吧
            ijkl, IJKL = self.vote_record(
                ijk[..., valid],
                I123[..., valid],
                th[..., valid],
                index[..., valid],
                n=i.shape[0],
            )
            if len(ijkl) > 4:
                # 成功一次后，直接用当前四点扩展
                # success, result = self.extend_identify(
                #     params, uncertainty, ijkl, IJKL, *args, factor=factor, **kwargs
                # )
                # if success:
                return ijkl, IJKL
            fail_num += 1
            if fail_num > max_try:
                break
        return None

    def vote_record(self, ijk, I123, th, index, n, vote_th=6):
        """
        投票法确定真实的匹配结果
        Arguments:
            params (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y, 1)
        Returns:
            list : matched ellipse ID
        """
        votes = torch.zeros(
            (self.catalog.shape[0], self.catalog.shape[0]),
            dtype=torch.int32,
            device=self.device,
        )
        # idx = np.arange(4, dtype=np.int32)
        for it in range(ijk.shape[2]):
            ABC = self.triad_search(I123[..., it], index[..., it], th[..., it])
            if ABC is None:
                continue
            pyra = ijk[..., it]
            assert np.unique(pyra[2, 1:]).shape[0] == 1
            ijkl = np.array((*pyra[:, 0], pyra[2, 1]))
            IJKL = self.build_pyramids(ABC)
            # 完成一次全图扩展
            if len(IJKL) > 0:
                votes[ijkl[None].repeat(IJKL.shape[0], axis=0), IJKL] += 1
        # 选取最大值
        votes = votes.cpu().numpy()
        ## 用匈牙利算法完成识别分配
        ijkl, IJKL = linear_sum_assignment(-votes)
        candidate = votes[ijkl, IJKL]
        ind = np.argsort(-candidate)[:8]
        ind = ind[candidate[ind] > 0]
        if ind.shape[0] < 4:
            return None
        else:
            return ijkl[ind], IJKL[ind]

    def solve_cross_points(self, m, Q):
        """求解直线m与椭圆Q的交点，以齐次形式返回"""
        assert Q.shape[0] == m.shape[0]
        # 用斜率判断非常不准确，改用角度
        res = m[:, 0].abs() < 1e6 * m[:, 1].abs()
        ## 角度也非常不准确，改为用相对的斜率
        cross = torch.zeros((3, 2, m.shape[0]), device=self.device)
        ## # 对于一般情况，用通用方法求解交点
        t_n = m[res, 2] / m[res, 1]
        m_n = m[res, 0] / m[res, 1]
        a = Q[res, 0, 0] - 2 * Q[res, 0, 1] * m_n + Q[res, 1, 1] * m_n**2
        b = 2 * (
            -Q[res, 0, 1] * t_n
            + Q[res, 1, 1] * m_n * t_n
            + Q[res, 0, 2]
            - Q[res, 1, 2] * m_n
        )
        c = Q[res, 1, 1] * t_n**2 - 2 * Q[res, 1, 2] * t_n + Q[res, 2, 2]
        ## 待解矩阵
        x = torch.stack(
            (
                (-b + (b**2 - 4 * a * c).sqrt()) / 2 / a,
                (-b - (b**2 - 4 * a * c).sqrt()) / 2 / a,
            )
        ).to(self.device)
        y = -m_n * x - t_n
        cross[..., res] = torch.stack((x, y, torch.ones_like(x).to(self.device)))
        # if (np.isclose(x[0], x[1]) & np.isclose(y[0], y[1])).sum() > 0:
        #     print("error")
        ## # 对于特殊情况，直线斜率为0，直接求解
        if not res.all():
            ind = ~res
            x = -m[ind, 2] / m[ind, 0]
            delta = (
                (Q[ind, 1, 2] + Q[ind, 0, 1] * x) ** 2
                - Q[ind, 1, 1]
                * (Q[ind, 2, 2] + 2 * Q[ind, 0, 2] * x + Q[ind, 0, 0] * x**2)
            ).sqrt()
            b = Q[ind, 1, 2] + Q[ind, 0, 1] * x
            x = torch.stack((x, x))
            y = torch.stack(
                (
                    (-b + delta) / Q[ind, 1, 1],
                    (-b - delta) / Q[ind, 1, 1],
                )
            )
            cross[..., ind] = torch.stack((x, y, torch.ones_like(x).to(self.device)))
        return cross.permute(2, 1, 0)

    def judge_inside(self, B, C, B_cross, C_cross):
        # 判断在内侧
        B = B[..., None, :]
        C = C[..., None, :]
        BB_cross = B_cross - B
        CB_cross = B_cross - C
        CC_cross = C_cross - C
        BC_cross = C_cross - B
        ## 判断是否在内侧，用小阈值排除等于零的情况
        index_B = (CB_cross * BB_cross).sum(dim=3) < -1e-3
        index_C = (CC_cross * BC_cross).sum(dim=3) < -1e-3
        # ind_B和ind_C应当是相等的，因为它们的原始顺序是一致的
        ind = (index_B.sum(dim=2) == 1).all(dim=0) & (index_C.sum(dim=2) == 1).all(
            dim=0
        )
        index_B = index_B & ind[None, :, None]
        index_C = index_C & ind[None, :, None]
        ## 返回所得交点和公共交集
        result = torch.ones_like(B_cross) * torch.nan
        result[:, ind] = torch.stack((B_cross[index_B], C_cross[index_C]), dim=1).view(
            3, -1, 2, 3
        )
        return result

    def g_func(self, x):
        assert x.shape[-2] == 6 and x.shape[-1] == 3
        # 获取五个点的组合
        comb = torch.combinations(torch.arange(6), 5).to(self.device)
        mu, nu = self.mu_nu(x[..., comb, :])
        # 计算不变量
        return torch.concatenate((mu, nu), dim=-1)

    def J(self, x):
        return (
            2 * x**6 - 6 * x**5 + 9 * x**4 - 8 * x**3 + 9 * x**2 - 6 * x + 2
        ) / (x**6 - 3 * x**5 + 3 * x**4 - x**3 + 3 * x**2 - 3 * x + 1)

    def mu_nu(self, x):
        mu = (
            torch.det(x[..., [3, 2, 0], :])
            * torch.det(x[..., [4, 1, 0], :])
            / torch.det(x[..., [3, 1, 0], :])
            / torch.det(x[..., [4, 2, 0], :])
        )
        nu = (
            torch.det(x[..., [3, 1, 0], :])
            * torch.det(x[..., [4, 2, 1], :])
            / torch.det(x[..., [3, 2, 1], :])
            / torch.det(x[..., [4, 1, 0], :])
        )
        return mu, nu
